import asyncio
import logging
from pathlib import PurePath, PureWindowsPath
from typing import List

from aardwolf import logger as aardwolf_logger
from aardwolf.commons.factory import RDPConnectionFactory
from aardwolf.commons.iosettings import RDPIOSettings
from aardwolf.commons.queuedata.constants import VIDEO_FORMAT
from aardwolf.connection import RDPConnection
from aardwolf.extensions.RDPECLIP.clipboard import Clipboard
from aardwolf.keyboard import VK_MODIFIERS
from aardwolf.keyboard.layoutmanager import KeyboardLayoutManager
from asyauth.common.constants import asyauthProtocol, asyauthSecret
from egg_timer import EggTimer

from common.credentials import Credentials, NTHash, Password, Username, get_plaintext
from common.types import NetworkPort
from infection_monkey.i_puppet import TargetHost

from .in_memory_file_provider import InMemoryFileProvider

logger = logging.getLogger(__name__)
aardwolf_logger.setLevel(logging.DEBUG)
RDP_HANDSHAKE = "S3Cr3TtH1Ng"
RESTART_FILE_EXPLORER_CMD = "cmd.exe /c taskkill /f /IM explorer.exe & start explorer.exe"
VK_RETURN = "VK_RETURN"


class RDPClient:
    def __init__(self):
        self._connection: RDPConnection = None
        self._authenticated = False
        self._kb_layout = KeyboardLayoutManager().get_layout_by_shortname("enus")
        self._file_provider = InMemoryFileProvider()
        self._event_loop = asyncio.SelectorEventLoop()
        self._save_screenshots = False
        asyncio.set_event_loop(self._event_loop)

    def connect(
        self, host: TargetHost, credentials: Credentials, port: NetworkPort, timeout: float
    ):
        _connect_task = self._event_loop.create_task(
            self._connect_with_timeout(host, credentials, port, timeout)
        )
        self._event_loop.run_until_complete(_connect_task)

    async def _connect_with_timeout(
        self, host: TargetHost, credentials: Credentials, port: NetworkPort, timeout: float
    ):
        try:
            await asyncio.wait_for(self._connect(host, credentials, port), timeout=timeout)
        except asyncio.TimeoutError:
            logger.debug(f"Failed to connect via RDP after {timeout} seconds")
            raise

    async def _connect(self, host: TargetHost, credentials: Credentials, port: NetworkPort):
        username = self._get_username(credentials)
        self._connection = self._create_rdp_connection(host, username)
        # Some chars are not supported in the URL, so we reset the password/nt_hash here
        self._connection.credentials.secret = self._get_secret(credentials)
        if isinstance(credentials.secret, NTHash):
            self._connection.credentials.protocol = asyauthProtocol.NTLM
            self._connection.credentials.stype = asyauthSecret.NT

        logger.debug(f"Connecting with user: {username}")
        _, err = await self._connection.connect()
        if err is not None:
            logger.exception(f"Error connecting via RDP: {err}")
            raise err

        logger.debug("RDP connection established")
        self._authenticated = True

    def _get_username(self, credentials: Credentials) -> str:
        if not isinstance(credentials.identity, Username):
            message = f"Unrecognized credential identity type: {credentials.identity}"
            logger.debug(message)
            raise ValueError(message)

        return credentials.identity.username

    def _get_secret(self, credentials: Credentials):
        if isinstance(credentials.secret, NTHash):
            return get_plaintext(credentials.secret.nt_hash)
        elif isinstance(credentials.secret, Password):
            return get_plaintext(credentials.secret.password)
        else:
            message = f"Unrecognized credential secret type: {credentials.secret}"
            logger.debug(message)
            raise ValueError(message)

    def _create_rdp_connection(self, host: TargetHost, username: str) -> RDPConnection:
        clipboard = Clipboard(file_provider=self._file_provider)
        target_settings = RDPClient._setup_iosettings(clipboard)
        url = f"rdp+ntlm-password://{username}:pass@{str(host.ip)}"
        connection_factory = RDPConnectionFactory.from_url(url, target_settings)

        return connection_factory.create_connection_newtarget(str(host.ip), target_settings)

    def copy_file(
        self,
        file: bytes,
        dest: PurePath,
        login_timeout: float,
        file_upload_timeout: float,
    ):
        try:
            _wait_for_screen_task = self._event_loop.create_task(
                self._wait_for_screen_to_load(login_timeout)
            )
            self._event_loop.run_until_complete(_wait_for_screen_task)
            is_screen_up = _wait_for_screen_task.result()

            if not is_screen_up:
                raise Exception(f"Desktop screen did not load in {login_timeout} seconds")

            _copy_file_task = self._event_loop.create_task(
                self._copy_file(file, dest, file_upload_timeout)
            )
            self._event_loop.run_until_complete(_copy_file_task)
        except asyncio.TimeoutError:
            logger.error(f"Failed to copy file to victim after {file_upload_timeout} seconds")
            raise

    async def _wait_for_screen_to_load(self, timeout: float):
        # On a fresh login, the screen may not be up yet. Wait for the desktop to load
        logger.debug("Waiting for the screen to load")
        if not await self._wait_for_machine_screen(timeout):
            return False

        # It is possible that a modal dialog, such as the permission prompt, is open.
        # In order to get rid of a modal dialog, we restart the file explorer
        logger.debug("Restarting file explorer")
        await self._remote_restart_file_explorer()

        # Restarting the file explorer may take extra time on startup, so we must again wait for
        # the screen to load
        return await self._wait_for_machine_screen(timeout)

    async def _copy_file(self, file: bytes, dest: PurePath, timeout: float):
        logger.debug(f"Opening a file explorer window with destination path: {dest}")
        await self._remote_open_folder(dest.parent)

        # If we don't have permission to access the folder, we will get a prompt
        # Press "Enter" to move past the prompt
        await self._send_vk_keypress(VK_RETURN)

        self._take_screenshot("05_before_copy.png")
        await self._copy_file_to_clipboard(dest.name, file)

        logger.debug(
            "Pasting the file from the clipboard to file explorer window with "
            f"destination path: {dest}"
        )
        await self._remote_paste_file(timeout)
        self._take_screenshot("06_after_paste.png")

        if not await self._remote_file_exists(dest):
            raise Exception(f"Failed to copy file to {dest}")

    async def _wait_for_machine_screen(self, timeout: float) -> bool:
        timer = EggTimer()
        timer.set(timeout)
        while not timer.is_expired():
            await asyncio.sleep(1)

            await self._open_windows_run_dialog()
            await self._send_keys(RDP_HANDSHAKE)
            await self._copy_selected_text_to_clipboard()
            remote_clipboard = await self._connection.get_current_clipboard_text()
            await self._synchronize_clipboard()
            if remote_clipboard == RDP_HANDSHAKE:
                return True
        self._take_screenshot("01_screen_timeout.png")

        return False

    async def _synchronize_clipboard(self):
        # Yield time for the clipboard to sync
        # Because the clipboard channel is async, we need to suspend our work to allow it to sync
        await asyncio.sleep(0.5)

    async def _remote_restart_file_explorer(self):
        await self._connection.set_current_clipboard_text(RESTART_FILE_EXPLORER_CMD)
        await self._synchronize_clipboard()
        await self._open_windows_run_dialog()
        await self._remote_paste()
        self._take_screenshot("02_restart_file_explorer.png")
        await self._send_vk_keypress(VK_RETURN)
        await asyncio.sleep(1)  # Wait for file explorer to reload

    def _take_screenshot(self, filename):
        if self._save_screenshots and self._connection.desktop_buffer_has_data:
            buffer = self._connection.get_desktop_buffer(VIDEO_FORMAT.PIL)
            buffer.save(filename)

    async def _remote_open_folder(self, path: PurePath):
        path_str = str(PureWindowsPath(path))
        await self._open_windows_run_dialog()
        await self._send_keys(path_str)
        self._take_screenshot("03_open_folder.png")
        await self._send_vk_keypress(VK_RETURN)
        # Wait for window to appear
        await asyncio.sleep(5)

        logger.debug(f"Opened file explorer window to: {path_str}")
        self._take_screenshot("04_folder_opened.png")

    async def _remote_file_exists(self, path: PurePath) -> bool:
        path_str = str(PureWindowsPath(path))
        verification_command = f'cmd /c dir "{path_str}" && echo true | clip || echo false | clip'
        await self._open_windows_run_dialog()
        await self._connection.set_current_clipboard_text(verification_command)
        await self._synchronize_clipboard()
        self._take_screenshot("07_file_exists.png")
        await self._remote_paste()
        await self._send_vk_keypress(VK_RETURN)
        result = await self._connection.get_current_clipboard_text()

        async with asyncio.timeout(5):
            while result == verification_command:
                await asyncio.sleep(0.5)
                result = await self._connection.get_current_clipboard_text()

        if result.strip() == "true":
            return True

        return False

    async def _copy_file_to_clipboard(self, filename: str, file: bytes):
        self._file_provider.set_files({filename: file})
        await self._connection.set_current_clipboard_files([filename])
        await self._synchronize_clipboard()
        logger.debug(f"File '{filename}' copied to clipboard")

    async def _remote_paste_file(self, timeout: float):
        async with asyncio.timeout(timeout):
            await self._remote_paste()
            while not self._file_provider.is_download_active():
                await self._remote_paste()
            while self._file_provider.is_download_active():
                await asyncio.sleep(1)

    async def _remote_paste(self):
        # Send Ctrl+V
        scancode, _ = self._kb_layout.char_to_scancode("v")
        modifier_scancode = self._kb_layout.vk_to_scancode("VK_LCONTROL")
        await self._send_keypress([modifier_scancode, scancode])
        await self._synchronize_clipboard()

    async def _copy_selected_text_to_clipboard(self):
        # Send Ctrl+A, Ctrl+C
        scancode, _ = self._kb_layout.char_to_scancode("a")
        modifier_scancode = self._kb_layout.vk_to_scancode("VK_LCONTROL")
        await self._send_keypress([modifier_scancode, scancode])

        scancode, _ = self._kb_layout.char_to_scancode("c")
        modifier_scancode = self._kb_layout.vk_to_scancode("VK_LCONTROL")
        await self._send_keypress([modifier_scancode, scancode])
        await self._synchronize_clipboard()

    @staticmethod
    def _setup_iosettings(clipboard: Clipboard) -> RDPIOSettings:
        target_settings = RDPIOSettings()
        target_settings.video_width = 1024
        target_settings.video_height = 768
        target_settings.video_bpp_max = 32
        target_settings.video_out_format = VIDEO_FORMAT.PIL
        target_settings.clipboard_use_pyperclip = False
        target_settings.clipboard = clipboard

        return target_settings

    def connected(self) -> bool:
        return self._authenticated

    def execute_command(self, command: str):
        try:
            _execute_task = self._event_loop.create_task(self._execute_and_verify_command(command))
            self._event_loop.run_until_complete(_execute_task)
        except Exception as err:
            logger.exception(f"Something happened while executing command: {err}")
            raise err

    async def _execute_and_verify_command(self, command: str):
        verifiable_command = f"{command} && echo true | clip || echo false | clip"
        await self._open_windows_run_dialog()
        await self._execute_command_in_shell(verifiable_command)
        result = await self._connection.get_current_clipboard_text()

        async with asyncio.timeout(10):
            while result == verifiable_command:
                await asyncio.sleep(0.5)
                result = await self._connection.get_current_clipboard_text()

        if result.strip() == "true":
            logger.debug("Command succeeded")
        elif result.strip() == "false":
            raise Exception("Command failed to run")
        else:
            raise Exception(f"Command returned unexpected result: '{result}'")

    async def _execute_command(self, command: str):
        await self._connection.set_current_clipboard_text(command)
        await self._synchronize_clipboard()
        await self._remote_paste()
        await self._send_vk_keypress(VK_RETURN)

    async def _execute_command_in_shell(self, command: str):
        logger.debug("Executing command in shell...")

        await self._execute_command("cmd.exe")
        await asyncio.sleep(1)

        await self._connection.set_current_clipboard_text(command)
        await self._synchronize_clipboard()

        # Because Copy/Paste using keyboard shortcuts is disabled by default in some
        # versions of Windows, we use ALT+SPACE,EP which uses the Explorer Window menu
        # to paste the command to Command Prompt
        alt = self._kb_layout.vk_to_scancode("VK_LMENU")
        space = self._kb_layout.vk_to_scancode("VK_SPACE")
        await self._send_keypress([alt, space])
        await asyncio.sleep(0.5)
        await self._send_keys("ep")

        self._take_screenshot("08_command_pasted.png")

        await self._send_vk_keypress(VK_RETURN)

    async def _open_windows_run_dialog(self):
        # Open run dialog
        scancode, _ = self._kb_layout.char_to_scancode("r")
        modifier_scancode = self._kb_layout.vk_to_scancode("VK_LWIN")
        await self._send_keypress([modifier_scancode, scancode])
        await asyncio.sleep(0.5)

    async def _send_keypress(self, scancodes: List[int]):
        if len(scancodes) == 0:
            return

        scancode = scancodes[0]
        await self._connection.send_key_scancode(
            scancode,
            is_pressed=True,
            is_extended=False,
        )
        await self._send_keypress(scancodes[1:])
        await self._connection.send_key_scancode(
            scancode,
            is_pressed=False,
            is_extended=False,
        )

    async def _send_vk_keypress(self, virtual_key: str):
        scancode = self._kb_layout.vk_to_scancode(virtual_key)
        await self._send_keypress([scancode])
        await asyncio.sleep(0.1)

    async def _send_keys(self, string: str):
        # Type the string
        for char in string:
            if char not in self._kb_layout.char_to_sc:
                logger.warning(f"Character '{char}' not found in layout {self._kb_layout.name}")
                continue

            scancode, modifiers = self._kb_layout.char_to_scancode(char)
            if modifiers == VK_MODIFIERS.VK_SHIFT or char.isupper():
                modifier_scancode = self._kb_layout.vk_to_scancode("VK_LSHIFT")
                await self._send_keypress([modifier_scancode, scancode])
            else:
                await self._send_keypress([scancode])

        await asyncio.sleep(0.2)

    def __del__(self):
        if self._connection is not None and not self._event_loop.is_closed():
            try:
                logger.debug("Terminating RDP connection")
                self._event_loop.run_until_complete(self._connection.terminate())
            except Exception:
                logger.debug("Failed to terminate RDP connection!")
        else:
            logger.debug("RDP connection already terminated")

        logger.debug("Canceling async event loop tasks")
        pending = asyncio.all_tasks(self._event_loop)
        for task in pending:
            task.cancel()

        logger.debug("Stopping async event loop")

        # Run loop until tasks done:
        self._event_loop.run_until_complete(asyncio.gather(*pending))

        logger.debug("Cleanup successful")
